package zk.learn.streaming

import java.util.Properties

import org.apache.kafka.clients.consumer.{ConsumerConfig, ConsumerRecord}
import org.apache.kafka.clients.producer.{KafkaProducer, ProducerConfig, ProducerRecord}
import org.apache.kafka.common.serialization.{StringDeserializer, StringSerializer}
import org.apache.spark.SparkConf
import org.apache.log4j.{Level, Logger}
import org.apache.spark.streaming.dstream.InputDStream
import org.apache.spark.streaming.kafka010.{ConsumerStrategies, KafkaUtils, LocationStrategies}
import org.apache.spark.streaming.{Seconds, StreamingContext}


object KafkaDStream {
  def main(args: Array[String]): Unit = {
    // 初始化spark
    Logger.getLogger("org").setLevel(Level.ERROR)
    val conf = new SparkConf().setAppName("KafkaDStream").setMaster("local[*]")
    val ssc = new StreamingContext(conf, Seconds(5))

    // 定义kafka相关参数
    val groupId: String = "mygroup01"
    val topics: Array[String] = Array("tp_spark_stream_input")
    val kafkaParams: Map[String, Object] = getKafkaConsumerParameters(groupId)

    // 创建DStream
    val dstream: InputDStream[ConsumerRecord[String, String]] = KafkaUtils.createDirectStream(
      ssc,
      LocationStrategies.PreferConsistent,
      ConsumerStrategies.Subscribe[String, String](topics, kafkaParams)
    )

    // 处理消息
    dstream.map(record => {
      // 切割
      record
        .value()
        .split(",")
        .map(_.replaceAll("<<<!>>>", ""))
        .mkString(" | ")
    }).foreachRDD{ (rdd, _) => {
      // 输出到kafka中
      rdd.foreachPartition(saveAsKafka)
    }}



    // 启动作业
    ssc.start()
    ssc.awaitTermination()
  }


  def saveAsKafka(iter: Iterator[String]): Unit = {
    // 构建生产者
    val brokers = "linux121:9092,linux122:9092,linux123:9092"
    val tpInput = "tp_spark_stream_output"
    val prop = new Properties()
    prop.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, brokers)
    prop.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, classOf[StringSerializer])
    prop.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, classOf[StringSerializer])
    val producer = new KafkaProducer[String, String](prop)

    // 输出到主题tp_spark_stream_output中
    iter.foreach(record => {
      val msg = new ProducerRecord[String, String](tpInput, record)
      producer.send(msg)
    })

  }

  def getKafkaConsumerParameters(groupid: String): Map[String, Object] = {
    Map[String, Object](
      ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG -> "linux121:9092,linux122:9092,linux123:9092",
      ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG -> classOf[StringDeserializer],
      ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG -> classOf[StringDeserializer],
      ConsumerConfig.GROUP_ID_CONFIG -> groupid,
      ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG -> (false: java.lang.Boolean),
      ConsumerConfig.AUTO_OFFSET_RESET_CONFIG -> "earliest"
    )
  }
}